from langchain.prompts.chat import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from pydantic import BaseModel, Field
from server.common import chat_model


def create_reminder_generator():
    """This function will generate and return a cron expression, which accepts user's reminder task."""
    llm = chat_model.qwen
    structured_llm_generator = llm.with_structured_output(GenerateReminder)
    system = """You are a reminder plan creator who can accurately understand the reminder task input by user and convert 
    it into formatted plan. A standard reminder plan should include a standardized cron expression and a reminder text. The tone of the 
    reminder text should be gentle and amiable. Attention please, cron expression must consist of 6 fields and reminder text must be Chinese."""
    generate_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system),
            (
                "human",
                "Here is the reminder task: \n\n {task} \n Formulate a reminder plan."
            )
        ]
    )
    return generate_prompt | structured_llm_generator


def create_reminder_grader():
    """This function will create a reminder grader, which accepts a reminder task and a cron expression.
    Then a result will be returned, 'yes' or 'no'. 'yes' denotes given is a right cron expression and correctly expresses the reminder task given."""
    llm = chat_model.qwen
    structured_llm_generator = llm.with_structured_output(ReminderGrader)
    system = """You are a grader assessing whether a cron expression generated by an LLM which is standardized and correctly expresses a reminder task. 
         Choose a score from 1 to 10. The higher the score, the better the expression."""
    generate_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system),
            (
                "human",
                "Here is the reminder task: \n\n {task} \n Here is the cron expression: \n\n {cron_expression}\n"
            )
        ]
    )
    return generate_prompt | structured_llm_generator


def create_task_rewriter():
    """This function will create a task rewriter, which accepts user's reminder task and returns a better version of
        it."""
    llm = chat_model.qwen
    structured_llm_generator = llm.with_structured_output(TaskRewriter)
    system = """You a reminder task re-writer that converts an input task to a better version that is optimized \n 
             for the generation of cron expression. Look at the input and try to reason about the underlying semantic intent / meaning."""
    re_write_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system),
            (
                "human",
                "Here is the initial reminder task: \n\n {task} \n Formulate an improved reminder task.",
            ),
        ]
    )
    return re_write_prompt | structured_llm_generator

class LLMS:
    """provide LLM for reminder_maker"""
    def __init__(self):
        self.reminder_generator = create_reminder_generator()
        self.reminder_grader = create_reminder_grader()
        self.task_rewriter = create_task_rewriter()


class GenerateReminder(BaseModel):
    """Formatted plan of user's reminder task."""

    cron_expression: str = Field(description="Cron expression of the reminder plan.")
    reminder_text: str = Field(description="Reminder text.")

class ReminderGrader(BaseModel):
    """Score of the cron expression which correctly expresses the reminder task, from 1 to 10."""

    binary_score: int = Field(description="The cron expression correctly expresses the reminder task, from 1 to 10.")

class TaskRewriter(BaseModel):
    """optimized reminder task."""

    reminder_task: str = Field(description="The new version of reminder task.")